{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h2q27gKz1H20"
      },
      "source": [
        "##### Copyright 2023 The MediaPipe Authors. All Rights Reserved."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TUfAcER1oUS6"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L_cQX8dWu4Dv"
      },
      "source": [
        "# Pose Landmarks Detection with MediaPipe Tasks\n",
        "\n",
        "This notebook shows you how to use MediaPipe Tasks Python API to detect pose landmarks from images."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6PN9FvIx614"
      },
      "source": [
        "## Preparation\n",
        "\n",
        "Let's start with installing MediaPipe.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gxbHBsF-8Y_l"
      },
      "outputs": [],
      "source": [
        "!pip install -q mediapipe"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a49D7h4TVmru"
      },
      "source": [
        "Then download an off-the-shelf model bundle. Check out the [MediaPipe documentation](https://developers.google.com/mediapipe/solutions/vision/pose_landmarker#models) for more information about this model bundle."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OMjuVQiDYJKF"
      },
      "outputs": [],
      "source": [
        "!wget -O pose_landmarker.task -q https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/1/pose_landmarker_heavy.task"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YYKAJ5nDU8-I"
      },
      "source": [
        "## Visualization utilities"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s3E6NFV-00Qt"
      },
      "outputs": [],
      "source": [
        "#@markdown To better demonstrate the Pose Landmarker API, we have created a set of visualization tools that will be used in this colab. These will draw the landmarks on a detect person, as well as the expected connections between those markers.\n",
        "\n",
        "from mediapipe import solutions\n",
        "from mediapipe.framework.formats import landmark_pb2\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "def draw_landmarks_on_image(rgb_image, detection_result):\n",
        "  pose_landmarks_list = detection_result.pose_landmarks\n",
        "  annotated_image = np.copy(rgb_image)\n",
        "\n",
        "  # Loop through the detected poses to visualize.\n",
        "  for idx in range(len(pose_landmarks_list)):\n",
        "    pose_landmarks = pose_landmarks_list[idx]\n",
        "\n",
        "    # Draw the pose landmarks.\n",
        "    pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()\n",
        "    pose_landmarks_proto.landmark.extend([\n",
        "      landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in pose_landmarks\n",
        "    ])\n",
        "    solutions.drawing_utils.draw_landmarks(\n",
        "      annotated_image,\n",
        "      pose_landmarks_proto,\n",
        "      solutions.pose.POSE_CONNECTIONS,\n",
        "      solutions.drawing_styles.get_default_pose_landmarks_style())\n",
        "  return annotated_image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83PEJNp9yPBU"
      },
      "source": [
        "## Download test image\n",
        "\n",
        "To demonstrate the Pose Landmarker API, you can download a sample image using the follow code. The image is from [Pixabay](https://pixabay.com/photos/girl-woman-fitness-beautiful-smile-4051811/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tzXuqyIBlXer"
      },
      "outputs": [],
      "source": [
        "!wget -q -O image.jpg https://cdn.pixabay.com/photo/2019/03/12/20/39/girl-4051811_960_720.jpg\n",
        "\n",
        "import cv2\n",
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "img = cv2.imread(\"image.jpg\")\n",
        "cv2_imshow(img)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u-skLwMBmMN_"
      },
      "source": [
        "Optionally, you can upload your own image. If you want to do so, uncomment and run the cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "etBjSdwImQPw"
      },
      "outputs": [],
      "source": [
        "# from google.colab import files\n",
        "# uploaded = files.upload()\n",
        "\n",
        "# for filename in uploaded:\n",
        "#   content = uploaded[filename]\n",
        "#   with open(filename, 'wb') as f:\n",
        "#     f.write(content)\n",
        "\n",
        "# if len(uploaded.keys()):\n",
        "#   IMAGE_FILE = next(iter(uploaded))\n",
        "#   print('Uploaded file:', IMAGE_FILE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iy4r2_ePylIa"
      },
      "source": [
        "## Running inference and visualizing the results\n",
        "\n",
        "The final step is to run pose landmark detection on your selected image. This involves creating your PoseLandmarker object, loading your image, running detection, and finally, the optional step of displaying the image with visualizations.\n",
        "\n",
        "Check out the [MediaPipe documentation](https://developers.google.com/mediapipe/solutions/vision/pose_landmarker/python) to learn more about configuration options that this solution supports.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_JVO3rvPD4RN"
      },
      "outputs": [],
      "source": [
        "# STEP 1: Import the necessary modules.\n",
        "import mediapipe as mp\n",
        "from mediapipe.tasks import python\n",
        "from mediapipe.tasks.python import vision\n",
        "\n",
        "# STEP 2: Create an PoseLandmarker object.\n",
        "base_options = python.BaseOptions(model_asset_path='pose_landmarker.task')\n",
        "options = vision.PoseLandmarkerOptions(\n",
        "    base_options=base_options,\n",
        "    output_segmentation_masks=True)\n",
        "detector = vision.PoseLandmarker.create_from_options(options)\n",
        "\n",
        "# STEP 3: Load the input image.\n",
        "image = mp.Image.create_from_file(\"image.jpg\")\n",
        "\n",
        "# STEP 4: Detect pose landmarks from the input image.\n",
        "detection_result = detector.detect(image)\n",
        "\n",
        "# STEP 5: Process the detection result. In this case, visualize it.\n",
        "annotated_image = draw_landmarks_on_image(image.numpy_view(), detection_result)\n",
        "cv2_imshow(cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_BwzFvaxwtPX"
      },
      "source": [
        "Visualize the pose segmentation mask."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3jAIFzw9M3JJ"
      },
      "outputs": [],
      "source": [
        "segmentation_mask = detection_result.segmentation_masks[0].numpy_view()\n",
        "visualized_mask = np.repeat(segmentation_mask[:, :, np.newaxis], 3, axis=2) * 255\n",
        "cv2_imshow(visualized_mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QipRi2ozw7cg"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
